import matplotlib
import os
import matplotlib.pyplot as plt
import numpy as np

def readmix(fname):
    infile = open(fname,'r')
    temp = infile.readline()

    while not 'NATOM' in temp:
        temp = infile.readline()
    temp = infile.readline().split() # vals
    natom, nmol, nreact, nfix, nvary = int(temp[0]), int(temp[1]), int(temp[2]), int(temp[4]), int(temp[6])
    temp = infile.readline() # header
    temp = infile.readline().split() # vals 
    nz = int(temp[2])

    # skip all lines that come before 'MIXING RATIO'
    temp = infile.readline()
    while 'MIXING RATIO' not in temp:
        temp = infile.readline()

    # loop over groups of species                   
    mix = {}
    nz = 60
    alt = np.zeros(nz)

    blank = infile.readline()
    header = infile.readline().split()
    alt = np.zeros(nz)
    while len(header)>1:
        for ix in range(1,len(header)):
            mix[header[ix]]= np.zeros(nz)
        for iz in range(0,nz):
            temp = infile.readline().split()
            for ix in range(1,len(header)):
                mix[header[ix]][iz] = float(temp[ix+1])
            alt[iz] = float(temp[1])
        blank = infile.readline()
        header = infile.readline().split()
    mix['z'] = alt
    return mix

mypath = ['outputs_1x','outputs_3x','outputs_10x']
dat = {}
for i in mypath:
    dat[i] = {}
    for ifile in os.listdir('./withsep/'+i+'/'):
        if not ifile.endswith('mini'):
            continue
        dat[i][ifile] = readmix('./withsep/'+i+'/'+ifile)

matplotlib.rc('font', family='serif')
matplotlib.rc('xtick', labelsize = 18)
matplotlib.rc('ytick', labelsize = 18)
matplotlib.rc('axes', linewidth=2)
fontname = 'serif'
fontsize=30
lw = 3.3


pre = ['7mbar','50mbar','200mbar','500mbar','1bar']
#species = ['NO','HNO2','HNO3','HO2NO2','NO2','N2O','N2O5']
#species = ['O2','O3','O','OH','HO2']
#species = ['N','NO','NO2','N2O']
#species = ['N','NO','O2','HO2']
#species = ['N','NO','NO2','N2O']
species = ['HNO2','HNO3','HO2NO2','N2O5']
clr = ['royalblue','darkorange','mediumvioletred','grey']
#clr = ['grey','royalblue','darkorange','mediumvioletred','turquoise','maroon','olivedrab']
#'green','purple','brown','orange','cyan','grey']

lns = ['-','--',':']
fig,ax=plt.subplots(ncols=5,nrows=1,figsize=(30,8))

for i in mypath:
    if  '1x' in i:
        ictr = 0
    elif '3x' in i:
        ictr = 1
    elif '10x' in i:
        ictr = 2
    else:
        print('ice case not found')
    for ifile in list(dat[i].keys()):
        kctr = 0

        if pre[0] in ifile:
            jctr = 0
        elif pre[1] in ifile:
            jctr = 1
        elif pre[2] in ifile:
            jctr = 2
        elif pre[3] in ifile:
            jctr = 3
        elif pre[4] in ifile:
            jctr = 4
        else:
            print('pressure not found, ifile')

        for isp in list(species):
            print(ictr,jctr,i,ifile)
            ax[jctr].semilogx(dat[i][ifile][isp],dat[i][ifile]['z'],color=clr[kctr],linestyle=lns[ictr],linewidth=2)
            kctr += 1
        ax[jctr].set_xlim([1e-17,1e-10])
        jctr += 1
    ictr += 1

print(list(dat['outputs_1x'].keys()))
print(list(dat['outputs_3x'].keys()))
print(list(dat['outputs_10x'].keys()))
print(list(dat.keys()))

ax[0].legend()
plt.savefig('mixing-bottleneck.png',dpi=300)

